# -*- coding: utf-8 -*-

import os
import sys
from PyPDF2 import PdfReader, PdfWriter
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
from reportlab.lib import colors
import io


def add_background_color(input_pdf_path, output_pdf_path, bg_color=(232, 246, 232)):
    """
    给PDF添加护眼背景色

    参数:
    input_pdf_path -- 输入PDF文件路径
    output_pdf_path -- 输出PDF文件路径
    bg_color -- 背景色RGB值，默认为淡绿色(232, 246, 232)，这是常用的护眼色
    """
    try:
        # 读取原始PDF
        reader = PdfReader(input_pdf_path)
        writer = PdfWriter()

        # 处理每一页
        for page_num in range(len(reader.pages)):
            # 获取原始页面
            page = reader.pages[page_num]
            page_width = float(page.mediabox.width)
            page_height = float(page.mediabox.height)

            # 创建一个带有背景色的新页面
            packet = io.BytesIO()
            c = canvas.Canvas(packet, pagesize=(page_width, page_height))

            # 设置背景色
            r, g, b = bg_color
            c.setFillColorRGB(r / 255, g / 255, b / 255)
            c.rect(0, 0, page_width, page_height, fill=True, stroke=False)
            c.save()

            # 将背景添加到页面下方
            packet.seek(0)
            background_pdf = PdfReader(packet)
            background_page = background_pdf.pages[0]

            # 合并背景和原始页面
            background_page.merge_page(page)
            writer.add_page(background_page)

        # 保存新PDF
        with open(output_pdf_path, "wb") as output_file:
            writer.write(output_file)

        print(f"已成功将护眼背景色添加到 {output_pdf_path}")
        return True
    except Exception as e:
        print(f"处理PDF时出错: {e}")
        return False


def main(input_pdf):
    if not os.path.exists(input_pdf):
        print(f"错误: 文件 '{input_pdf}' 不存在!")
        return

    # 生成输出文件名（原文件名_护眼版.pdf）
    file_name, file_ext = os.path.splitext(input_pdf)
    output_pdf = f"{file_name}_护眼版{file_ext}"

    # 询问用户是否要自定义背景色
    custom_color = input("是否要自定义背景色? (y/n，默认为淡绿色护眼色): ").lower()

    if custom_color == "y":
        try:
            r = int(input("请输入红色值 (0-255): "))
            g = int(input("请输入绿色值 (0-255): "))
            b = int(input("请输入蓝色值 (0-255): "))
            bg_color = (r, g, b)
        except ValueError:
            print("输入无效，将使用默认护眼色")
            bg_color = (232, 246, 232)
    else:
        bg_color = (232, 246, 232)  # 默认护眼色

    # 添加背景色
    success = add_background_color(input_pdf, output_pdf, bg_color)

    if success:
        print(f"处理完成！新文件已保存为: {output_pdf}")
    else:
        print("处理失败，请检查文件和权限。")


if __name__ == "__main__":
    main(rf"xxx.pdf")
